#include "Binding/BoundTreeRewriter.hpp"

#include "Binding/BoundAssignment.hpp"
#include "Binding/BoundBinary.hpp"
#include "Binding/BoundBlock.hpp"
#include "Binding/BoundConditionalGotoStatement.hpp"
#include "Binding/BoundExpressionStatement.hpp"
#include "Binding/BoundForStatement.hpp"
#include "Binding/BoundIfStatement.hpp"
#include "Binding/BoundNode.hpp"
#include "Binding/BoundUnary.hpp"
#include "Binding/BoundVariableDeclaration.hpp"
#include "Binding/BoundWhileStatement.hpp"

#include <cstddef>
#include <memory>
#include <vector>

SharedBoundStatement BoundTreeRewriter::RewriteStatement(SharedBoundStatement node)
{
    switch (node->GetKind())
    {
    case BoundNodeKind::BlockStatement :
        return this->RewriteBlockStatement(node);
    case BoundNodeKind::VariableDeclaration :
        return this->RewriteVariableDeclaration(node);
    case BoundNodeKind::IfStatement :
        return this->RewriteIfStatement(node);
    case BoundNodeKind::WhileStatement :
        return this->RewriteWhileStatement(node);
    case BoundNodeKind::ForStatement :
        return this->RewriteForStatement(node);
    case BoundNodeKind::LabelStatement :
        return this->RewriteLabelStatement(node);
    case BoundNodeKind::GotoStatement :
        return this->RewriteGotoStatement(node);
    case BoundNodeKind::ConditionalGotoStatement :
        return this->RewriteConditionalGotoStatement(node);
    case BoundNodeKind::ExpressionStatement :
        return this->RewriteExpressionStatement(node);
    default :
        throw "Unexpected node:" + node->GetKindName();
    }
}

SharedBoundStatement BoundTreeRewriter::RewriteBlockStatement(SharedBoundStatement node)
{
    auto                              blockNode = (BoundBlockStatement *)node.get();
    std::vector<SharedBoundStatement> builder;

    for (size_t i = 0; i < blockNode->Statements.size(); i++)
    {
        auto oldStatement = blockNode->Statements[i];
        auto newStatement = this->RewriteStatement(oldStatement);
        if (oldStatement != newStatement)
        {
            if (builder.empty())
            {
                for (size_t j = 0; j < i; j++)
                    builder.push_back(blockNode->Statements[j]);
            }
        }
        if (!builder.empty())
            builder.push_back(newStatement);
    }
    if (builder.empty())
    {
        return node;
    }
    return std::make_shared<BoundBlockStatement>(builder);
}

SharedBoundStatement BoundTreeRewriter::RewriteVariableDeclaration(SharedBoundStatement node)
{
    auto variable    = (BoundVariableDeclaration *)node.get();
    auto initializer = this->RewriteExpression(variable->Initializer);
    if (initializer == variable->Initializer)
        return node;
    return std::make_shared<BoundVariableDeclaration>(variable->Variable, initializer);
}

SharedBoundStatement BoundTreeRewriter::RewriteIfStatement(SharedBoundStatement node)
{
    auto ifNode        = (BoundIfStatement *)node.get();
    auto condtion      = this->RewriteExpression(ifNode->Condition);
    auto thanStatement = this->RewriteStatement(ifNode->ThenStatement);
    auto elseStatement = ifNode->ElseStatement == nullptr ? nullptr : this->RewriteStatement(ifNode->ElseStatement);

    if ((condtion == ifNode->Condition) && (thanStatement == ifNode->ThenStatement) && (elseStatement == ifNode->ElseStatement))
        return node;

    return std::make_shared<BoundIfStatement>(condtion, thanStatement, elseStatement);
}

SharedBoundStatement BoundTreeRewriter::RewriteWhileStatement(SharedBoundStatement node)
{
    BoundWhileStatement *whileNode = (BoundWhileStatement *)node.get();

    SharedBoundExpression condtion = this->RewriteExpression(whileNode->Condition);
    SharedBoundStatement  body     = this->RewriteStatement(whileNode->BodyStatement);

    if ((condtion == whileNode->Condition) && (body == whileNode->BodyStatement))
        return node;

    return std::make_shared<BoundWhileStatement>(condtion, body);
}
SharedBoundStatement BoundTreeRewriter::RewriteForStatement(SharedBoundStatement node)
{
    BoundForStatement *forNode       = (BoundForStatement *)node.get();

    SharedBoundExpression lowerBound = this->RewriteExpression(forNode->LowerBound);
    SharedBoundExpression upperBound = this->RewriteExpression(forNode->UpperBound);
    SharedBoundStatement  body       = this->RewriteStatement(forNode->Body);

    if ((lowerBound == forNode->LowerBound) && (upperBound == forNode->UpperBound) && (body == forNode->Body))
        return node;

    return std::make_shared<BoundForStatement>(forNode->Variable, lowerBound, upperBound, body);
}
SharedBoundStatement BoundTreeRewriter::RewriteExpressionStatement(SharedBoundStatement node)
{
    BoundExpressionStatement *es     = (BoundExpressionStatement *)node.get();

    SharedBoundExpression expression = this->RewriteExpression(es->Expression);

    if (expression == es->Expression)
        return node;

    return std::make_shared<BoundExpressionStatement>(expression);
}
SharedBoundStatement BoundTreeRewriter::RewriteLabelStatement(SharedBoundStatement node)
{
    return node;
}
SharedBoundStatement BoundTreeRewriter::RewriteGotoStatement(SharedBoundStatement node)
{
    return node;
}

SharedBoundExpression BoundTreeRewriter::RewriteExpression(SharedBoundExpression node)
{
    switch (node->GetKind())
    {
    case BoundNodeKind::ErrorExpression :
        return RewriteErrorExpression(node);
    case BoundNodeKind::LiteralExpression :
        return RewriteLiteralExpression(node);
    case BoundNodeKind::VariableExpression :
        return RewriteVariableExpression(node);
    case BoundNodeKind::AssignmentExpression :
        return RewriteAssignmentExpression(node);
    case BoundNodeKind::UnaryExpression :
        return RewriteUnaryExpression(node);
    case BoundNodeKind::BinaryExpression :
        return RewriteBinaryExpression(node);
    default :
        throw "Unexpected node:" + node->GetKindName();
    }
}
SharedBoundExpression BoundTreeRewriter::RewriteErrorExpression(SharedBoundExpression node)
{
    return node;
}
SharedBoundExpression BoundTreeRewriter::RewriteLiteralExpression(SharedBoundExpression node)
{
    return node;
}
SharedBoundExpression BoundTreeRewriter::RewriteVariableExpression(SharedBoundExpression node)
{
    return node;
}

SharedBoundStatement BoundTreeRewriter::RewriteConditionalGotoStatement(SharedBoundStatement node)
{
    auto cgotoNode = (BoundConditionalGotoStatement *)node.get();

    auto condtion  = this->RewriteExpression(cgotoNode->Condtion);
    if (condtion == cgotoNode->Condtion)
    {
        return node;
    }
    return std::make_shared<BoundConditionalGotoStatement>(cgotoNode->Lable, condtion, cgotoNode->JumpIfTrue);
}
SharedBoundExpression BoundTreeRewriter::RewriteAssignmentExpression(SharedBoundExpression node)
{
    BoundAssignmentExpression *exprNode = (BoundAssignmentExpression *)node.get();

    SharedBoundExpression expression    = this->RewriteExpression(exprNode->Expression);
    if (expression == exprNode->Expression)
        return node;

    return std::make_shared<BoundAssignmentExpression>(exprNode->Variable, expression);
}
SharedBoundExpression BoundTreeRewriter::RewriteUnaryExpression(SharedBoundExpression node)
{
    BoundUnaryExpression *exprNode   = (BoundUnaryExpression *)node.get();

    SharedBoundExpression expression = this->RewriteExpression(exprNode->Operand);
    if (expression == exprNode->Operand)
        return node;

    return std::make_shared<BoundUnaryExpression>(exprNode->Op, expression);
}
SharedBoundExpression BoundTreeRewriter::RewriteBinaryExpression(SharedBoundExpression node)
{
    BoundBinaryExpression *exprNode = (BoundBinaryExpression *)node.get();

    SharedBoundExpression left      = this->RewriteExpression(exprNode->Left);
    SharedBoundExpression right     = this->RewriteExpression(exprNode->Right);
    if ((left == exprNode->Left) && (right == exprNode->Right))
        return node;

    return std::make_shared<BoundBinaryExpression>(left, exprNode->Op, right);
}
